from starlette.applications import Starlette
from uvicorn import Config, Server
import asyncio
from mcp.server.auth.routes import create_auth_routes
from provider import LeohOAuthServerProvider
from pydantic import AnyHttpUrl
from mcp.server.auth.settings import ClientRegistrationOptions, RevocationOptions
from starlette.requests import Request
from starlette.responses import Response, HTMLResponse, RedirectResponse, JSONResponse
from starlette.routing import Route
from jinja2 import Environment, FileSystemLoader
from starlette.exceptions import HTTPException
from mcp.server.auth.provider import AuthorizationCode, construct_redirect_uri, AccessToken
from mcp.server.auth.routes import cors_middleware
import time

jinja2_env = Environment(loader=FileSystemLoader('templates'))


host = "localhost"
port = 9090

oauth_provider = LeohOAuthServerProvider(
    login_url=f"http://{host}:{port}/login",
)


async def login_handler(request: Request) -> Response:
    state = request.query_params.get("state")
    html = jinja2_env.get_template("login.html").render(
        {
            "login_callback_url": "/login/callback",
            "state":state,
        }
    )
    return HTMLResponse(html)

async def login_callback_handler(request: Request) -> Response:
    # 验证用户名和密码，如果验证成功，就要返回授权码
    form = await request.form()
    state = form.get("state")
    username = form.get("username")
    password = form.get("password")

    if not username or not password or not state:
        raise HTTPException(status_code=400, detail="Missing username, password or state parameters")

    if not isinstance(username, str) or not isinstance(password, str) or not isinstance(state, str):
        raise HTTPException(status_code=400, detail="Invalid parameters types")

    # 验证用户名和密码
    if username != "demo_user" and password != "demo_password":
        raise HTTPException(status_code=401, detail="Invalid credentials")

    # 创建授权码
    auth_code : AuthorizationCode = await oauth_provider.create_authorization_code(state)
    redirect_uri = auth_code.redirect_uri
    # 重定向到redirect_uri，因为重定向的时候，需要携带参数：state、code（授权码）
    # http://localhost:3030/callback?code=xxx&state=xxx
    return RedirectResponse(construct_redirect_uri(str(redirect_uri), code=auth_code.code, state=state), status_code=302)


async def introspect_handler(request: Request) -> Response:
    form = await request.form()
    token = form.get("token")
    if not token:
        return JSONResponse({"active": False})

    access_token = await oauth_provider.load_access_token(token)
    if not access_token:
        return JSONResponse({"active": False})
    if access_token.expires_at < time.time():
        return JSONResponse({"active": False})

    return JSONResponse({
        "active": True,
        "client_id": access_token.client_id,
        "scope": " ".join(access_token.scopes),
        "exp": access_token.expires_at,
        "iat": int(time.time()),
        "token_type": "Bearer",
        "aud": access_token.resource,
    })


routes = [
    Route("/login", login_handler, methods=['GET']),
    Route("/login/callback", login_callback_handler, methods=['POST']),
    Route("/introspect", cors_middleware(introspect_handler, ['POST', 'OPTIONS']), methods=['POST', 'OPTIONS']),
]

routes += create_auth_routes(
    provider=oauth_provider,
    issuer_url=AnyHttpUrl(f"http://{host}:{port}"),
    # 开启客户端注册的功能
    client_registration_options=ClientRegistrationOptions(
        enabled=True,
        valid_scopes=['user'],
        default_scopes=['user'],
    ),
    revocation_options=RevocationOptions(
        enabled=True,
    )
)

async def main():
    app = Starlette(debug=True, routes=routes)
    config = Config(
        app=app,
        host=host,
        port=port,
        log_level="info",
    )
    server = Server(config)
    await server.serve()

if __name__ == '__main__':
    asyncio.run(main())